/* 
 * File:   dns_query.h
 * Author: liujianping@snda.com
 * NOTE: use this class, when compile you need link libresolv
 * Created on 2011年3月23日, 上午9:32
 */

#ifndef DNS_QUERY_H
#define	DNS_QUERY_H

#include <netinet/in.h>
#include <arpa/nameser.h>
#include <resolv.h>
#include "log.h"

#include <string>
#include <vector>
using namespace std;

class dns_query_t
{
public:
    static int mx_query(const char* domain, vector<string> & mx_list, vector<string>& default_mx);
};


unsigned short getshort(unsigned char *c) { return (c[0] << 8) + c[1]; }

int dns_query_t::mx_query(const char* domain, vector<string> & mx_list, vector<string>& default_mx)
{
    logtrace((COMMON, "dns_query_t::mx_query begin with domain:<%s> ...", domain));
    // make a clean start
    if(default_mx != mx_list)
    {
        mx_list.clear();
        mx_list.assign(default_mx.begin(), default_mx.end());
    }
    
    // initialize resolver
    if (res_init()) {
        logerror((COMMON, "dns_query_t::mx_query => res_init() failed"));
        return -1;
    }

    // do MX query
    union {
        HEADER hdr;
        u_char buf[NS_PACKETSZ];
    } response;

    int response_len = res_query(domain, C_IN, T_MX, (u_char *)&response, sizeof(response));
    if (response_len < 0) {
        logerror((COMMON, "dns_query_t::mx_query => res_query() failed"));
        return -1;
    }
    if ((size_t)response_len > sizeof(response)) {
        response_len = sizeof(response);
    }

    u_char * response_end = response.buf + response_len;
    u_char * response_pos = response.buf + sizeof(HEADER);

    // skip queries
    int response_count = ntohs(response.hdr.qdcount);
    while (response_count-- > 0) {
        char query_name[MAXDNAME];
        int bytes_read = dn_expand(response.buf, response_end, response_pos, query_name, MAXDNAME);
        if (bytes_read < 0) {
            logerror((COMMON, "dns_query_t::mx_query => dn_expand() failed:parse QUERY part in MX lookup 1."));
            return -1;
        }
        response_pos += bytes_read;
        if ((response_end - response_pos) < QFIXEDSZ) {
            logerror((COMMON, "dns_query_t::mx_query => dn_expand() failed:parse QUERY part in MX lookup 2."));
            return -1;
        }
        response_pos += QFIXEDSZ;
    }

    // now process answers
    response_count = ntohs(response.hdr.ancount);
    while (response_count-- > 0) {
        char mx_name[MAXDNAME];
        int bytes_processed = dn_expand(response.buf, response_end, response_pos, mx_name, MAXDNAME);
        if (bytes_processed < 0) {
            logerror((COMMON, "dns_query_t::mx_query => dn_expand() failed:parse QUERY part in MX lookup 3."));
            return -1;
        }

        response_pos += bytes_processed;
        if ((response_end - response_pos) < 10) {
            logerror((COMMON, "dns_query_t::mx_query => dn_expand() failed:parse QUERY part in MX lookup 4."));
            return -1;
        }

        unsigned short rrtype = getshort(response_pos);
        unsigned short rrdlen = getshort(response_pos + 8);
        response_pos += 10;
        if (rrtype == T_MX) {
            //!unsigned short priority = (response_pos[0] << 8) + response_pos[1];
            memset(mx_name, 0, MAXDNAME);
            if (dn_expand(response.buf, response_end, response_pos+2, mx_name, MAXDNAME) < 0) {
                logerror((COMMON, "dns_query_t::mx_query => dn_expand() failed:parse QUERY part in MX lookup 5."));
                return -1;
            }
            //! VERBOSE("MTA::get_mx_list(): entry [" << mx_name << "] priority [" << priority << "]");
            mx_list.push_back(std::string(mx_name));
        }
        response_pos += rrdlen;
    }
    
    logtrace((COMMON, "dns_query_t::mx_query domain:<%s> end ok.", domain));
    return mx_list.empty();
}

#endif	/* DNS_QUERY_H */

